
import sys
import indexing.featureindex
from theconfig import *

resComp = { "A":"T", "T":"A", "C":"G", "G":"C", "N":"N", "a":"t", "t":"a", "g":"c", "c":"g", "n":"n" }

def toStrand(strand, base):
    if strand=="+": return base
    elif strand=="-": return resComp[base]
    raise Exception("bad strand: " + strand)

def readCodonTable(filename=CODON_TABLE):
    shorthands = { '*' : 'ACGT', '(' : 'CT', ')' : 'AG' }
    table = {}
    f = open(filename, "r")
    for line in f:
        cols = line.strip().split("\t")
        if line.startswith("#") or len(cols)!=3:
            continue
        (codon, aa, aaChar) = tuple(cols)
        if len(codon)!=3:
            continue
        if shorthands.has_key(codon[-1]):
            codons = [codon[:2] + ch for ch in shorthands[codon[-1]]]
        else:
            codons = [codon]
        for codon in codons:
            table[codon] = (aa, aaChar)
    f.close()
    return table

def readBlosumMatrix(codonTable, filename=BLOSUM_MATRIX):
    codons = {}
    for (pattern, (aa, aaChar)) in codonTable.items():
        codons[aaChar] = aa
    rows = []
    mat = {}
    for line in open(filename):
        if line.startswith("#"):
            continue
        entries = filter(lambda(x):x!='', line.strip().split(" "))
        if not rows:
            rows = [codons.get(entry, entry) for entry in entries]
            # print >> sys.stderr, "# Read BLOSUM column headers: ", rows
            for row in rows:
                mat[row] = {}
            continue
        if len(entries) != len(rows)+1:
            # print >> sys.stderr, "# Ignoring BLOSUM file line: ", line.strip()
            continue
        row = codons.get(entries[0], entries[0])
        for ix in xrange(1, len(entries)):
            mat[row][rows[ix-1]] = int(entries[ix])
    # print >> sys.stderr, "# BLOSUM matrix: ", mat
    return mat

def getDbsnp(query):
    (qchr, qmin, qmax) = query
    ret = {}
    for feature in indexing.featureindex.getFeatures(DBSNP_FILE, (qchr, qmin, qmax)):
        ret[(feature.chr, feature.spos)] = True
    return ret

def readDbsnpSnps(filename=DBSNP_FILE, dbsnp={}, chromosomes=None):
    print >> sys.stderr, "# Reading SNPs from file", filename, "into map of size", len(dbsnp)
    indexed = filename.endswith(".INDEX")
    if indexed and chromosomes and len(chromosomes)!=1:
        for chr in chromosomes.keys():
            readDbsnpSnps(filename, dbsnp, { chr : True })
        return
    if indexed:
        chr = list(chromosomes.keys())[0]
        features = indexing.featureindex.getFeatures(filename, (chr, None, None))
        print >> sys.stderr, "# Reading indexed features from chromosome ", chr
    else:
        features = open(filename, "r")
        print >> sys.stderr, "# Reading unindexed features from file ", filename
    numlines = 0
    for feature in features:
        numlines += 1
        if 0 == numlines%100000:
            print >> sys.stderr, "# Read %d lines from %s" % (numlines, filename)
        if indexed:
            line = feature.content
        else:
            line = feature
        cols = line.strip().split("\t")
        if len(cols)!=6 and len(cols)!=3:
            print >> sys.stderr, "# Warning: ignoring line", line.strip()
            continue
        if chromosomes and not chromosomes.has_key(cols[0]):
            # ignoring this chromosome
            continue
        if len(cols)==6:
            # dbSNP131/132 format
            key = (cols[0], long(cols[1]))
            if not dbsnp.has_key(cols[0]):
                dbsnp[cols[0]] = {}
            if not dbsnp[cols[0]].has_key(long(cols[1])):
                dbsnp[cols[0]][long(cols[1])] = []
            dbsnp[cols[0]][long(cols[1])].append("%s:%s:%s:val=%s" % (cols[5], cols[2], cols[3], cols[4]))
        else:
            key = (cols[0], cols[1], cols[2])
            dbsnp[key] = True
    if not indexed:
        features.close()
    print >> sys.stderr, "# Done... read %d SNPs from %s" % (len(dbsnp), filename)
    return dbsnp

def getDbsnpByLocus(dbsnps, chr, pos):
    return dbsnps.get(chr, {}).get(pos, None)

def getDbsnpByMutation(dbsnps, chr, pos, oldBase, newBase):
    muts = dbsnps.get(chr, {}).get(pos, [])
    if not muts: return None
    ret = []
    for mut in muts:
        (tmp1, alleles, strand, tmp2) = tuple(mut.split(":"))
        alleleList = alleles.split("/")
        if strand=="+" and not oldBase in alleleList and resComp[oldBase] in alleleList:
            strand = "-"
        if strand=="+":
            if (oldBase in alleleList and newBase in alleleList):
                ret.append(mut)
        elif strand=="-":
            if (resComp[oldBase] in alleleList and resComp[newBase] in alleleList):
                ret.append(mut)
    return ",".join(ret)

class GeneInfo:
    def __init__(self, chr, txStart, txEnd, geneId, geneSymbol, cdsStart, cdsEnd, strand, exonStarts, exonEnds):
        self.chr = chr
        self.txStart = txStart
        self.txEnd = txEnd
        self.geneId = geneId
        self.geneSymbol = geneSymbol
        self.cdsStart = long(cdsStart)
        self.cdsEnd = long(cdsEnd)
        self.strand = strand
        self.exonStarts = exonStarts
        self.exonEnds = exonEnds

    def containsPosition(self, chr, pos):
        return (chr==self.chr and pos>=self.txStart and pos<=self.txEnd)

    def getIntronicLocationDesc(self, chr, pos):
        starts = map(int, self.exonStarts.split(",")[:-1])
        ends = map(int, self.exonEnds.split(",")[:-1])
        for i in xrange(len(starts)):
            if pos>ends[i] and pos<(1+starts[i+1]):
                dist5pDonor = pos - ends[i]
                dist3pAcceptor = (1+starts[i+1]) - pos
                if self.strand == "-":
                    tmp = dist5pDonor
                    dist5pDonor = dist3pAcceptor
                    dist3pAcceptor = tmp
                if dist5pDonor <= dist3pAcceptor:
                    return "non-coding-intronic,%d-from-5p-donor" % dist5pDonor
                else:
                    return "non-coding-intronic,%d-from-3p-acceptor" % dist3pAcceptor
        return "non-coding-intronic,unknown-pos"
    
    def getNonCodingLocationDesc(self, chr, pos):
        if not self.containsPosition(chr, pos):
            return "intergenic"
        elif self.cdsStart == self.cdsEnd:
            return "non-coding-gene"
        elif (self.strand=="+" and pos<self.cdsStart) or (self.strand=="-" and pos>self.cdsEnd):
            return "non-coding-5-prime-utr"
        elif (self.strand=="+" and pos>self.cdsEnd) or (self.strand=="-" and pos<self.cdsStart):
            return "non-coding-3-prime-utr"
        else:
            return self.getIntronicLocationDesc(chr, pos)

class GeneInfoTable:
    def __init__(self, src):
        if src=="UCSC":
            geneInfoFile = UCSC_KNOWN_GENE_TABLE
            geneSymbolFile = UCSC_KG_XREF_TABLE
        else:
            geneInfoFile = CCDS_KNOWN_GENE_TABLE
            geneSymbolFile = CCDS_GENE_NAMES

        lookup = {}
        lookupById = {}
        symbols = {}

        f = open(geneSymbolFile, "r")
        for line in f:
            if line.startswith("#"): continue
            cols = line.strip().split("\t")
            geneId = cols[0]
            if src=="UCSC":
                # if using UCSC X-ref table
                if cols[4]!="":
                    symbols[geneId] = cols[4]
            else:
                # if using CCDS names table
                symbols[geneId] = cols[1]
        f.close()

        f = open(geneInfoFile, "r")
        for line in f:
            if line.startswith("#"): continue
            cols = line.strip().split("\t")
            if src=="UCSC":
                cols = cols[:10]
            else:
                cols = cols[1:11]
            (geneId, chr, strand, txStart, txEnd, cdsStart, cdsEnd, exonCount, exonStarts, exonEnds) = tuple(cols)
            if not lookup.has_key(chr):
                lookup[chr] = []
            gene = GeneInfo(chr, long(txStart), long(txEnd), geneId, symbols.get(geneId, None), long(cdsStart), long(cdsEnd), strand, exonStarts, exonEnds)
            lookup[chr].append(gene)
            if lookupById.has_key((chr, geneId)):
                raise Exception("multiple genes on chr " + chr + " with id " + geneId)
            lookupById[(chr, geneId)] = gene
        f.close()
        for key in lookup.keys():
            # lookup[key].sort()
            pass # print >> sys.stderr, "# %d entries in lookup table for %s" % (len(lookup[key]), key)
        self.genesByChr = lookup
        self.genesById = lookupById

    def getGenes(self, chr, pos):
        lookupChr = self.genesByChr[chr]
        genes = []
        for gene in lookupChr:
            if gene.containsPosition(chr, pos):
                genes.append(gene)
        return genes
    
    def getGeneSymbol(self, chr, pos):
        symbol = None
        genes = self.getGenes(chr, pos)
        for gene in genes:
            if gene.geneSymbol:
                symbol = gene.geneSymbol
                break
        if not symbol:
            for gene in genes:
                symbol = gene.geneId
                break
        if not symbol:
            symbol = "(None)"
        return symbol
    
    def getGeneIds(self, chr, pos):
        genes = self.getGenes(chr, pos)
        return [gene.geneId for gene in genes]
    
    def getNonCodingLocationDesc(self, chr, pos, geneId):
        if not self.genesById.has_key((chr, geneId)):
            raise Exception("couldn't find gene on chr " + chr + " with id " + geneId)
        gene = self.genesById[(chr, geneId)]
        return (gene.strand, gene.getNonCodingLocationDesc(chr, pos))

if __name__=='__main__':
    table = GeneInfoTable("/exds/bowtie/references/ccds_codon_info/ccdsGene.table.txt", "/exds/bowtie/references/ccds_codon_info/ccdsToGene.via-SureSelect.txt")
